open Ast
open Lexing
open Cub
open Atom
open Format

module StrMap = Map.Make(String)

(*
  TODO : compilation

*)

type error = Not_a_line

(* Exception generique utilisee au typage *)
exception Error of error

let line t = 
  match t with Tlinenum loc -> loc
  | _ -> raise (Error Not_a_line)
    
type kind = Sconst | Sglobal | Slocal | Sarray  

let procsloc = Hashtbl.create 7

let maxloc = ref 0
  
let numprocs = ref 0

let locs = ref []

let unsafes = ref []
  
let id_trans = ref 0
let next_id () = incr id_trans; !id_trans
let app_id n = Printf.sprintf "%s_%d" n (next_id ())

let app_loc loc = Printf.sprintf "L%d" loc
  
let getloc p =
  let loc = (try Hashtbl.find procsloc p with Not_found -> 0) + 1 in
  Hashtbl.add procsloc p loc;
  let l = app_loc loc in
  if loc > !maxloc then (incr maxloc; locs := l :: !locs);
  l

let peekloc p =
  app_loc ((try Hashtbl.find procsloc p with Not_found -> 0) + 1)

let mk_trans name args reqs ureq assigns arrass upds nondets =
  {
    tr_name = name; tr_args = args;
    tr_reqs = reqs; tr_ureq = ureq;
    tr_assigns = assigns; tr_arrass = arrass;
    tr_upds = upds; tr_nondets = nondets;
  }
    
let mk_trans_id name = mk_trans (app_id name)

let mk_node e t = {data = t; node = e}

let loc_dummy e =
  mk_node e (Lexing.dummy_pos, Lexing.dummy_pos)

let not_of_compop = function
| Bne -> Beq
| Bge -> Blt
| Ble -> Bgt
| Beq -> Bne
| Blt -> Bge
| Bgt -> Ble

let pid_of_ident s = String.sub s 1 (String.length s - 1)
    
let not_of_logicop = function
| Band -> Bor
| Bor -> Band

let rec not_of_lt lt =
  match lt.node with
  | LTterm t ->
      begin match t.node with
      | TMatom at ->
          begin match at.node with
          | ATbool b ->
              mk_node (
                LTterm
                  (mk_node (TMatom (mk_node (ATbool (not b)) Tbool)) Tbool) 
              ) Tbool
          | _ -> lt end
      | _ -> lt end
  | LTlit l ->
      let compop, t1, t2 = l.node in
      let l' = mk_node (not_of_compop compop, t1, t2) Tbool in
      mk_node (LTlit l') Tbool
  | LTop (logicop, lt1, lt2) -> mk_node
      (LTop (not_of_logicop logicop, not_of_lt lt1, not_of_lt lt2)) Tbool 

    
let type_string = function
| Tint -> "int"
| Treal -> "real"
| Tproc -> "proc"
| Tbool -> "bool"
| Tident id -> id.node
| _ -> failwith "Anomaly"

let multiply_tr trans dnf =
  match dnf with
  | QDforall (str, ssa) -> 
      [Strans { trans with
        tr_name = app_id trans.tr_name;
        tr_ureq = [str, SSAtom.elements ssa];}]
  | QDnone ssa -> 
      List.rev (SSAtom.fold (
        fun sa tl ->
          let tr = { trans with
            tr_name = app_id trans.tr_name;
            tr_reqs = SAtom.union sa trans.tr_reqs;
          } in
          (Strans tr) :: tl
      ) ssa [])
  | QDexists (str, ssa) -> 
      List.rev (SSAtom.fold (
        fun sa tl ->
          let tr = { trans with
            tr_name = app_id trans.tr_name;
            tr_reqs = SAtom.union sa trans.tr_reqs;
            tr_args = List.rev (str::trans.tr_args);
          } in
          (Strans tr) :: tl
      ) ssa [])

let multiply_us (p, g) dnf =
  let ssa, ps = match dnf with
  | QDnone ssa -> ssa, p
  | QDforall (_, ssa) -> assert false (* TODO *)
  | QDexists (str, ssa) -> ssa, str::p
  in
  List.rev (SSAtom.fold (
    fun sa ul ->
      let us = (ps, SAtom.union sa g) in
      (Sunsafe us) :: ul
  ) ssa [])

let req_sn p sn =
  SAtom.singleton
    (Comp ((Access ("PC", [p])), Eq, Const (ConstName sn)))

let rec replace_loc oldloc newloc states = 
  match states with
  | Sapp (sts1, sts2) -> Sapp ((replace_loc oldloc newloc sts1),
                               (replace_loc oldloc newloc sts2))
  | Slist stsl ->
      Slist (List.map ( fun st ->
        begin match st with
        | Strans tr ->
            let arrass =
              List.map (fun a ->
                let id, l, t = a in
                if id = "PC" then
                begin match t with
                | Elem s when s = oldloc -> id, l, Elem newloc
                | _ -> a end
                else a) tr.tr_arrass
            in Strans { tr with tr_arrass = arrass}
        | _ -> st
        end) stsl)

let tr_name ?proc:(p="") name line  =
  Printf.sprintf "%s_l%d%s" name line p
        
let rec compile_assign lv t = 
  let t = compile_term t in
  match lv.node with
  | LVident (id) ->
      if Typing.is_global id.node then
      id.node.[0] <- Char.uppercase id.node.[0];
      Avar (id.node, t)
  | LVaccess (id, idl) ->
      let idl = List.map (fun e ->
        if Typing.is_global e.node then
        e.node.[0] <- Char.uppercase e.node.[0];
        e.node) idl in
      Aarray (id.node, idl, t)
        
and compile_stmt proc extra_dnf stmt =
  let curprocs, pid = if !numprocs > 1
    then [], "_p" ^ pid_of_ident proc else [proc], "" in
  let bpos, epos = line stmt.data in
  match stmt.node with
  | Scritical ->
      let sn = getloc proc in
      let reqs = req_sn proc sn in
      let arrass = [("PC", [proc], (Elem (peekloc proc)))] in
      let tr = mk_trans_id
        (tr_name ~proc:pid "crit" bpos.pos_lnum)
        curprocs reqs [] [] arrass [] [] in
      if not (List.mem (proc, peekloc proc) !unsafes) then
      unsafes := (proc, peekloc proc) :: !unsafes;
      state_l [Strans tr]

  | Sassert qt ->
      let sn = peekloc proc in
      let dnf = compile_qterm qt in
      let reqs = req_sn proc sn in
      let us = (curprocs, reqs)  in
      state_l (multiply_us us dnf)

  | Sawait (qt, None) ->
      let sn = getloc proc in
      let dnf = compile_qterm qt in
      let reqs = req_sn proc sn in
      let arrass = [("PC", [proc], (Elem (peekloc proc)))] in
      let tr = mk_trans
        (tr_name ~proc:pid "await" bpos.pos_lnum)
        curprocs reqs [] []
        (List.rev arrass) [] [] in
      state_l (multiply_tr tr dnf)

  | Sawait (qt, Some s) ->
      let dnf = compile_qterm qt in
      compile_stmt proc (Some dnf) s

  | Sassign (lv, t) ->
      let sn = getloc proc in
      let reqs = req_sn proc sn in
      let assign = compile_assign lv t in
      let pc_assign = ( "PC", [proc], (Elem (peekloc proc))) in
      let arrass, varass = begin match assign with  
      | Avar (id,  t) -> [pc_assign], [(id, t)]
      | Aarray (id, idl, t) -> (id, idl, t) :: [pc_assign], []
      end in
      let tr = mk_trans_id
        (tr_name ~proc:pid "assign" bpos.pos_lnum)
        curprocs reqs [] varass arrass [] [] in
      begin match extra_dnf with
      | None -> state_l [Strans tr]
      | Some dnf -> state_l (multiply_tr tr dnf)
      end

  | Satomic stmts ->
      let sn = getloc proc in
      let reqs = req_sn proc sn in
      let assigns, arrass =
        List.fold_left (
          fun ((ass, arr) as acc) e ->
            begin match e.node with
            | Sassign (lv, t) ->
                let assign = compile_assign lv t in
                begin match assign with 
                | Avar (id,  t) -> (id, t) :: ass, arr
                | Aarray (id, idl, t) -> ass, (id, idl, t) :: arr
                end
            | _ -> acc
            end
        ) ([], []) stmts in
      let arrass = ( "PC", [proc], (Elem (peekloc proc))) :: arrass in
      let tr = mk_trans_id
        (tr_name ~proc:pid "atomic" bpos.pos_lnum)
        curprocs reqs [] (List.rev assigns)
        (List.rev arrass) [] [] in
      begin match extra_dnf with
      | None -> state_l [Strans tr]
      | Some dnf -> state_l (multiply_tr tr dnf)
      end

  | Swhile (qt, sl) ->
      let sn = getloc proc in
      let dnf = compile_qterm qt in
      let ndnf = compile_qterm (not_of_qt qt) in
      let reqs = req_sn proc sn in

      (* Transition pour while = true *)
      let arrass = [("PC", [proc], (Elem (peekloc proc)))] in
      let tr_wt = mk_trans
        (tr_name ~proc:pid "while_true" bpos.pos_lnum)
        curprocs reqs
        [] [] (List.rev arrass) [] [] in
      let tr_wt_l = state_l (multiply_tr tr_wt dnf) in

      (* Transition pour les statements *)
      let tr_stmt_l = compile_stmts proc sl in

      (* Transition pour boucler *)
      let reqs_l = req_sn proc (getloc proc) in
      let arrass = [("PC", [proc], (Elem (sn)))] in
      let tr_wl = mk_trans_id
        (tr_name ~proc:pid "while_loop" epos.pos_lnum)
        curprocs reqs_l
        [] [] (List.rev arrass) [] [] in
      let tr_wl_l = state_l [Strans tr_wl] in

      (* Transition pour while = false *)
      let arrass = [("PC", [proc], (Elem (peekloc proc)))] in
      let tr_wf = mk_trans
        (tr_name ~proc:pid "while_false" bpos.pos_lnum)
        curprocs reqs
        [] [] (List.rev arrass) [] [] in
      let tr_wf_l = state_l (multiply_tr tr_wf ndnf) in

      tr_wt_l +++ tr_stmt_l +++ tr_wl_l +++ tr_wf_l
  | Sifelse (qt, sl1, sl2) ->

      let dnf = compile_qterm qt in
      let ndnf = compile_qterm (not_of_qt qt) in
      let sn = getloc proc in
      let reqs = req_sn proc sn in
      (* on sauvegarde la loc *)
      let beg_stmt1 = peekloc proc in
      let tr_stmt1_l = compile_stmts proc sl1 in

      (* on reserve *)
      let je_sn = peekloc proc in

      (* on sauvegarde la loc *)
      let beg_stmt2 = peekloc proc in
      let tr_stmt2_l = compile_stmts proc sl2 in

      let tr_stmt1_l = replace_loc je_sn (peekloc proc) tr_stmt1_l in

      (* transition condition = true *)
      let arrass = [("PC", [proc], (Elem beg_stmt1))] in
      let tr_it = mk_trans
        (tr_name ~proc:pid "if_true" bpos.pos_lnum)
        curprocs reqs
        [] [] (List.rev arrass) [] [] in
      let tr_it_l = state_l (multiply_tr tr_it dnf) in
      
      (* transition condition = false *)
      let arrass = [("PC", [proc], (Elem beg_stmt2))] in
      let tr_if = mk_trans
        (tr_name ~proc:pid "if_false" bpos.pos_lnum)
        curprocs reqs
        [] [] (List.rev arrass) [] [] in
      let tr_if_l = state_l (multiply_tr tr_if ndnf) in
      
      tr_it_l +++ tr_if_l +++ tr_stmt1_l +++ tr_stmt2_l

and compile_stmts proc sl =
  let rec compile_stmts states sl = 
    match sl with
    | [] -> states
    | stmt::sl -> compile_stmts (states +++ compile_stmt proc None stmt) sl
  in
  compile_stmts nost sl

and compile_leftval lv =
  match lv.node with
  | LVident id ->
      if Typing.is_global id.node then
      id.node.[0] <- Char.uppercase id.node.[0];
      Elem id.node
  | LVaccess (id, idl) ->
      id.node.[0] <- Char.uppercase id.node.[0];
      Access (id.node, List.map (fun e ->
        if Typing.is_global e.node then
        e.node.[0] <- Char.uppercase e.node.[0];
        e.node) idl)

and compile_constnum c =
  match c.node with
  | Creal r -> ConstReal r
  | Cint i -> ConstInt i

and compile_atomic_term at =
  match at.node with
  | ATleftval lv -> compile_leftval lv
  | ATconstnum c -> Const (compile_constnum c)
  | ATbool b -> if b then Elem "True" else Elem "False"
  | ATenum id -> Elem id.node

and compile_term t =
  match t.node with
  | TMatom at -> compile_atomic_term at
  | TMarith (aop, at, c) ->
      let aop = begin match aop with
      | Bplus -> Plus
      | Bminus -> Minus
      end in
      let at = compile_atomic_term at in
      let c = compile_constnum c in
      Arith (at, aop, c)

and compile_literal l =
  let (op, t1, t2) = l.node in
  let ct1 = compile_term t1 in
  let ct2 = compile_term t2 in
  SAtom.singleton (
    begin match op with
    | Bne -> Comp (ct1, Neq, ct2)
    | Bge -> Comp (ct2, Lt, ct1)
    | Ble -> Comp (ct1, Le, ct2)
    | Beq -> Comp (ct1, Eq, ct2)
    | Blt -> Comp (ct1, Lt, ct2)
    | Bgt -> Comp (ct2, Le, ct1)
    end)
    
and distrib s1 s2 =
  SSAtom.fold (fun e1 acc1 ->
    SSAtom.fold (fun e2 acc2 ->
      let satom = SAtom.union e1 e2 in
      SSAtom.add satom acc2
    ) s2 acc1 
  ) s1 SSAtom.empty 

and purednf lt =
  match lt.node with
  | LTop (Band, lt1, lt2) ->
      distrib (purednf lt1) (purednf lt2) 
  | LTop (Bor, lt1, lt2) ->
      SSAtom.union (purednf lt1) (purednf lt2) 
  | LTterm t -> 
      SSAtom.singleton
        (SAtom.singleton (Comp (compile_term t, Eq, (Elem "True"))))
  | LTlit l -> SSAtom.singleton (compile_literal l)
      

and not_of_qt qt =
  match qt.node with
  | Qterm lt ->
      mk_node (Qterm (not_of_lt lt)) Tany
  | Qquantifier (quantif, ident, lt) ->
      let q = (match quantif with Qexists -> Qforall | Qforall -> Qexists) in
      mk_node (Qquantifier (q, ident, not_of_lt lt)) Tany
        
        
and compile_qterm qt =
  match qt.node with
  | Qterm lt ->
      QDnone (purednf lt)
  | Qquantifier (quantif, ident, lt) ->
      begin match quantif with
      | Qexists -> QDexists (ident.node, purednf lt)
      | Qforall -> QDforall (ident.node, purednf lt)
      end

and compile_type_decl (id, idl) =
  let ids = List.fold_left (fun ids id -> id.node :: ids) [] idl in
  (id.node, List.rev ids)

and compile_var_decl = function
| VDconst (id, t) ->
    let s = String.copy id.node in
    s.[0] <- Char.uppercase id.node.[0];
    Dconst (s, (type_string t))
| VDglobal (id, t) ->
    let s = String.copy id.node in
    s.[0] <- Char.uppercase id.node.[0];
    Dglobal (s, (type_string t))
| VDarray (id, it, t) ->
    Darray (id.node, ([type_string it], type_string t))

and compile_decls (dl : (Ast.typ Ast.decl) list) : elem list * decls =
  List.fold_left(
    fun (ctl, cvdl) d ->
      match d with
      | Dtype td -> ((compile_type_decl td) :: ctl, cvdl)
      | Dvar vd -> (ctl, cvdl ++ (decl_l [compile_var_decl vd]))
  ) ([],nop) dl

let make_us_exists p1 p2 loc1 loc2 =
  let aloc1 = (Comp ((Access ("PC", [p1])), Eq, Const (ConstName loc1))) in
  let aloc2 = (Comp ((Access ("PC", [p2])), Eq, Const (ConstName loc2))) in
  state_l [Sunsafe ([p1;p2], SAtom.add aloc2 (SAtom.singleton aloc1))]
    
let make_us p1 p2 loc1 loc2 =
  let aloc1 = (Comp ((Access ("PC", [p1])), Eq, Const (ConstName loc1))) in
  let aloc2 = (Comp ((Access ("PC", [p2])), Eq, Const (ConstName loc2))) in
  state_l [Sunsafe ([], SAtom.add aloc2 (SAtom.singleton aloc1))]

let rec gen_unsafes stl unsafes =
  match unsafes with 
  | [] -> stl
  | (proc, loc) :: unsafes ->
      let stl =
        (if !numprocs = 1 then make_us_exists "p1" "p2" loc loc else nost) +++
          List.fold_left (fun accst (proc', loc') -> 
            if !numprocs = 1 then begin
              accst +++ make_us_exists "p1" "p2" loc loc'
            end
            else if proc' = proc then accst
            else accst +++ make_us proc proc' loc loc'
          ) stl unsafes in
      gen_unsafes stl unsafes

        
let compile_prog (p : Ast.typ Ast.prog) : Cub.system =
  let decls, procs = p in
  numprocs := List.length procs;
  let ctdl, cvdl = compile_decls decls in
  let cp, (sa_init, stl) =
    match procs with
    | [] -> None, (SAtom.empty, nost)
    | [ident, stmts] -> 
        let aloc = (Comp ((Access ("PC", [ident.node])), Eq,
                          Const (ConstName "L1"))) in
        Some ident.node, (SAtom.singleton aloc,
                          compile_stmts ident.node stmts)
    | l -> None, List.fold_left
      (fun (sa, accst) (ident, stmts) ->
        let st = compile_stmts ident.node stmts in
        let aloc = (Comp ((Access ("PC", [ident.node])), Eq,
                          Const (ConstName "L1"))) in
        SAtom.add aloc sa, accst +++ st
      ) (SAtom.empty, nost) l
  in 
  let ctdl = ("location", List.rev ((app_loc (!maxloc + 1))::!locs)) :: ctdl in
  let cvdl = decl_l [Darray ("PC", (["proc"], "location"))] ++ cvdl in
  let stl = state_l [Sinit (cp, sa_init)] +++ stl in
  let stl = gen_unsafes stl !unsafes in
  {
    num_procs = !numprocs;
    type_defs = List.rev ctdl;
    decls = cvdl;
    states = stl;
  }










